package jm.data.structure.heap;

import jm.data.structure.tree.printer.BinaryTrees;

import java.util.Comparator;

/**
 * @Description 使用小顶堆实现 top k
 * @date 2022/4/29 20:37
 */
public class TopKHeap {
    public static void main(String[] args) {
        Integer[] nums = {29, 66, 67, 96, 43, 99, 64, 89, 78, 17, 100, 69, 12};
        int k = 5;
        getTopMaxK(nums, k);
        getTopMinK(nums, k);
    }

    /**
     * 获取最大的前 k 个元素
     * @param nums
     * @param k
     */
    public static void getTopMaxK(Integer[] nums, int k){
        BinaryHeap<Integer> heap = new BinaryHeap<>(new Comparator<Integer>() {
            @Override
            public int compare(Integer o1, Integer o2) {
                return o2 - o1;
            }
        });
        for (int i = 0; i < nums.length; i++) {
            if (heap.size() < k){
                heap.add(nums[i]);
            }else {
                if (nums[i] > heap.get()){ // 如果当前元素大于堆顶元素，则加入到堆中。
                    heap.replace(nums[i]);
                }
            }
        }
        BinaryTrees.println(heap);
    }

    /**
     * 获取最小的前 k 个元素
     * @param nums
     * @param k
     */
    public static void getTopMinK(Integer[] nums, int k){
        BinaryHeap<Integer> heap = new BinaryHeap<>(new Comparator<Integer>() {
            @Override
            public int compare(Integer o1, Integer o2) {
                return o1 - o2;
            }
        });
        for (int i = 0; i < nums.length; i++) {
            if (heap.size() < k){
                heap.add(nums[i]);
            }else {
                if (nums[i] < heap.get()){ // 如果当前元素小于堆顶元素，则加入到堆中。
                    heap.replace(nums[i]);
                }
            }
        }
        BinaryTrees.println(heap);
    }
}
